#include "StdAfx.h"
#include "NtDriverController.h"
#include "Common.h"
#include "SysUtil.h"

CNtDriverController::CNtDriverController(void):
		m_hSCM(NULL),
		m_hDriver(NULL),
		m_bDriverStarted(FALSE),
		m_bErrorOnStart(FALSE)
{
	if (TRUE == Open())
	{
		strcpy(m_szName, "NTProcDrv");
		strcpy(m_szInfo, "Process creation detector for NT.");
		char szFullFileName[MAX_PATH];
		GetProcessHostFullName(szFullFileName);
		if (TRUE == ReplaceFileName(szFullFileName, "NtProcDrv.sys", m_szFullFileName))
			m_bDriverStarted = InstallAndStart();
	} // if
}

CNtDriverController::~CNtDriverController(void)
{
	StopAndRemove();
	Close();
}

BOOL CNtDriverController::Open()
{
	m_hSCM = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
	return (NULL != m_hSCM);
}

void CNtDriverController::Close()
{
	if(NULL != m_hDriver)
	{
		::CloseServiceHandle(m_hDriver);
		m_hDriver=NULL;
	}
	if(NULL != m_hSCM)
	{
		::CloseServiceHandle(m_hSCM);
		m_hSCM=NULL;
	}
}
BOOL CNtDriverController::WaitForStatus (
	DWORD dwDesiredStatus,
	SERVICE_STATUS * pss
	)
{
	BOOL bResult=FALSE;

	if(NULL != m_hDriver)
	{
		while(1)
		{
			bResult = ::QueryServiceStatus(m_hDriver,pss);
			if(!bResult)
				break;
			if(pss->dwCurrentState == dwDesiredStatus)
				break;

			DWORD dwWaitHint = pss->dwWaitHint / 10;    // Poll 1/10 of the wait hint
			if (dwWaitHint <  1000) dwWaitHint = 1000;  // At most once a second
			if (dwWaitHint > 10000) dwWaitHint = 10000; // At least every 10 seconds
			::Sleep(dwWaitHint);
		}
	}
	return bResult;
}
BOOL CNtDriverController::InstallAndStart()
{
	BOOL bResult=FALSE;

	if(NULL != m_hSCM)
	{
		m_hDriver=::CreateService(
			m_hSCM, 
			m_szName, 
			m_szInfo,
			SERVICE_ALL_ACCESS,
			SERVICE_KERNEL_DRIVER,
			SERVICE_DEMAND_START,
			SERVICE_ERROR_NORMAL,
			m_szFullFileName, 
			NULL, 
			NULL,
			NULL, 
			NULL, 
			NULL
			);
		if (NULL == m_hDriver)
		{
			if(	(::GetLastError() == ERROR_SERVICE_EXISTS)||
				(::GetLastError() == ERROR_SERVICE_MARKED_FOR_DELETE) )
			{
				m_hDriver = ::OpenService(
					m_hSCM,
					m_szName,
					SERVICE_ALL_ACCESS
					);
			}
		}else
		{
			SERVICE_STATUS serviceStatus = { 0 };
			bResult = ::StartService(m_hDriver, 0, NULL);
			if (bResult)
				bResult = WaitForStatus(SERVICE_RUNNING, &serviceStatus);	
			else
				bResult = (::GetLastError() == ERROR_SERVICE_ALREADY_RUNNING);
			// We should call DeleteService() if the SCM reports an error
			// on StartService(). Otherwise, the service will remain loaded
			// in an undesired state
			if (!bResult)
			{
				// Mark the service for deletion.
				::DeleteService(m_hDriver);
				if (m_hDriver != NULL)
				{
					::CloseServiceHandle(m_hDriver);
					m_hDriver = NULL;
				}
				m_bErrorOnStart = TRUE;
			}
		}
		return bResult;
					
	}
}

void CNtDriverController::StopAndRemove()
{
	if ((NULL != m_hDriver) && (!m_bErrorOnStart))
	{
		BOOL bResult;
		SERVICE_STATUS serviceStatus = { 0 };
		// Notifies a service that it should stop. 
		bResult = ::ControlService(m_hDriver, SERVICE_CONTROL_STOP, &serviceStatus);
		if (bResult)
			bResult = WaitForStatus(SERVICE_STOPPED, &serviceStatus);	
		// Mark the service for deletion.
		::DeleteService(m_hDriver);
	} // if
}
